{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|default_exp models.InceptionTimePlus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# InceptionTimePlus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ">This is an unofficial PyTorch implementation of InceptionTime (Fawaz, 2019) created by Ignacio Oguiza."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**References:**\n",
    "* Fawaz, H. I., Lucas, B., Forestier, G., Pelletier, C., Schmidt, D. F., Weber, J., ... & Petitjean, F. (2020). [Inceptiontime: Finding alexnet for time series classification. Data Mining and Knowledge Discovery, 34(6), 1936-1962.](https://arxiv.org/pdf/1909.04939)\n",
    "* Official InceptionTime tensorflow implementation: https://github.com/hfawaz/InceptionTime"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "from tsai.imports import *\n",
    "from collections import OrderedDict\n",
    "from fastai.layers import *\n",
    "from tsai.utils import *\n",
    "from tsai.models.layers import *\n",
    "from tsai.models.utils import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "# This is an unofficial PyTorch implementation by Ignacio Oguiza - oguiza@timeseriesAI.co modified from:\n",
    "\n",
    "# Fawaz, H. I., Lucas, B., Forestier, G., Pelletier, C., Schmidt, D. F., Weber, J., ... & Petitjean, F. (2019). \n",
    "# InceptionTime: Finding AlexNet for Time Series Classification. arXiv preprint arXiv:1909.04939.\n",
    "# Official InceptionTime tensorflow implementation: https://github.com/hfawaz/InceptionTime\n",
    "\n",
    "    \n",
    "class InceptionModulePlus(Module):\n",
    "    def __init__(self, ni, nf, ks=40, bottleneck=True, padding='same', coord=False, separable=False, dilation=1, stride=1, conv_dropout=0., sa=False, se=None,\n",
    "                 norm='Batch', zero_norm=False, bn_1st=True, act=nn.ReLU, act_kwargs={}):\n",
    "        \n",
    "        if not (is_listy(ks) and len(ks) == 3):\n",
    "            if isinstance(ks, Integral): ks = [ks // (2**i) for i in range(3)]\n",
    "            ks = [ksi if ksi % 2 != 0 else ksi - 1 for ksi in ks]  # ensure odd ks for padding='same'\n",
    "\n",
    "        bottleneck = False if ni == nf else bottleneck\n",
    "        self.bottleneck = Conv(ni, nf, 1, coord=coord, bias=False) if bottleneck else noop # \n",
    "        self.convs = nn.ModuleList()\n",
    "        for i in range(len(ks)): self.convs.append(Conv(nf if bottleneck else ni, nf, ks[i], padding=padding, coord=coord, separable=separable,\n",
    "                                                         dilation=dilation**i, stride=stride, bias=False))\n",
    "        self.mp_conv = nn.Sequential(*[nn.MaxPool1d(3, stride=1, padding=1), Conv(ni, nf, 1, coord=coord, bias=False)])\n",
    "        self.concat = Concat()\n",
    "        self.norm = Norm(nf * 4, norm=norm, zero_norm=zero_norm)\n",
    "        self.conv_dropout = nn.Dropout(conv_dropout) if conv_dropout else noop\n",
    "        self.sa = SimpleSelfAttention(nf * 4) if sa else noop\n",
    "        self.act = act(**act_kwargs) if act else noop\n",
    "        self.se = nn.Sequential(SqueezeExciteBlock(nf * 4, reduction=se), BN1d(nf * 4)) if se else noop\n",
    "        \n",
    "        self._init_cnn(self)\n",
    "    \n",
    "    def _init_cnn(self, m):\n",
    "        if getattr(self, 'bias', None) is not None: nn.init.constant_(self.bias, 0)\n",
    "        if isinstance(self, (nn.Conv1d,nn.Conv2d,nn.Conv3d,nn.Linear)): nn.init.kaiming_normal_(self.weight)\n",
    "        for l in m.children(): self._init_cnn(l)\n",
    "\n",
    "    def forward(self, x):\n",
    "        input_tensor = x\n",
    "        x = self.bottleneck(x)\n",
    "        x = self.concat([l(x) for l in self.convs] + [self.mp_conv(input_tensor)])\n",
    "        x = self.norm(x)\n",
    "        x = self.conv_dropout(x)\n",
    "        x = self.sa(x)\n",
    "        x = self.act(x)\n",
    "        x = self.se(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "@delegates(InceptionModulePlus.__init__)\n",
    "class InceptionBlockPlus(Module):\n",
    "    def __init__(self, ni, nf, residual=True, depth=6, coord=False, norm='Batch', zero_norm=False, act=nn.ReLU, act_kwargs={}, sa=False, se=None, \n",
    "                 stoch_depth=1., **kwargs):\n",
    "        self.residual, self.depth = residual, depth\n",
    "        self.inception, self.shortcut, self.act = nn.ModuleList(), nn.ModuleList(), nn.ModuleList()\n",
    "        for d in range(depth):\n",
    "            self.inception.append(InceptionModulePlus(ni if d == 0 else nf * 4, nf, coord=coord, norm=norm, \n",
    "                                                      zero_norm=zero_norm if d % 3 == 2 else False,\n",
    "                                                      act=act if d % 3 != 2 else None, act_kwargs=act_kwargs, \n",
    "                                                      sa=sa if d % 3 == 2 else False,\n",
    "                                                      se=se if d % 3 != 2 else None,\n",
    "                                                      **kwargs))\n",
    "            if self.residual and d % 3 == 2:\n",
    "                n_in, n_out = ni if d == 2 else nf * 4, nf * 4\n",
    "                self.shortcut.append(Norm(n_in, norm=norm) if n_in == n_out else ConvBlock(n_in, n_out, 1, coord=coord, bias=False, norm=norm, act=None))\n",
    "                self.act.append(act(**act_kwargs))\n",
    "        self.add = Add()\n",
    "        if stoch_depth != 0: keep_prob = np.linspace(1, stoch_depth, depth)\n",
    "        else: keep_prob = np.array([1] * depth)\n",
    "        self.keep_prob = keep_prob\n",
    "\n",
    "    def forward(self, x):\n",
    "        res = x\n",
    "        for i in range(self.depth):\n",
    "            if self.keep_prob[i] > random.random() or not self.training:\n",
    "                x = self.inception[i](x)\n",
    "            if self.residual and i % 3 == 2: \n",
    "                res = x = self.act[i//3](self.add(x, self.shortcut[i//3](res)))\n",
    "        return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@delegates(InceptionModulePlus.__init__)\n",
    "class InceptionTimePlus(nn.Sequential):\n",
    "    def __init__(self, c_in, c_out, seq_len=None, nf=32, nb_filters=None,\n",
    "                 flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None, custom_head=None, **kwargs):\n",
    "        \n",
    "        if nb_filters is not None: nf = nb_filters\n",
    "        else: nf = ifnone(nf, nb_filters) # for compatibility\n",
    "        backbone = InceptionBlockPlus(c_in, nf, **kwargs)\n",
    "        \n",
    "        #head\n",
    "        self.head_nf = nf * 4\n",
    "        self.c_out = c_out\n",
    "        self.seq_len = seq_len\n",
    "        if custom_head is not None: \n",
    "            if isinstance(custom_head, nn.Module): head = custom_head\n",
    "            else: head = custom_head(self.head_nf, c_out, seq_len)\n",
    "        else: head = self.create_head(self.head_nf, c_out, seq_len, flatten=flatten, concat_pool=concat_pool, \n",
    "                                      fc_dropout=fc_dropout, bn=bn, y_range=y_range)\n",
    "            \n",
    "        layers = OrderedDict([('backbone', nn.Sequential(backbone)), ('head', nn.Sequential(head))])\n",
    "        super().__init__(layers)\n",
    "        \n",
    "    def create_head(self, nf, c_out, seq_len, flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None):\n",
    "        if flatten: \n",
    "            nf *= seq_len\n",
    "            layers = [Flatten()]\n",
    "        else: \n",
    "            if concat_pool: nf *= 2\n",
    "            layers = [GACP1d(1) if concat_pool else GAP1d(1)]\n",
    "        layers += [LinBnDrop(nf, c_out, bn=bn, p=fc_dropout)]\n",
    "        if y_range: layers += [SigmoidRange(*y_range)]\n",
    "        return nn.Sequential(*layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "class InCoordTime(InceptionTimePlus):\n",
    "    def __init__(self, *args, coord=True, zero_norm=True, **kwargs):\n",
    "        super().__init__(*args, coord=coord, zero_norm=zero_norm, **kwargs)\n",
    "\n",
    "\n",
    "class XCoordTime(InceptionTimePlus):\n",
    "    def __init__(self, *args, coord=True, separable=True, zero_norm=True, **kwargs):\n",
    "        super().__init__(*args, coord=coord, separable=separable, zero_norm=zero_norm, **kwargs)\n",
    "        \n",
    "InceptionTimePlus17x17 = named_partial('InceptionTimePlus17x17', InceptionTimePlus, nf=17, depth=3)\n",
    "InceptionTimePlus32x32 = named_partial('InceptionTimePlus32x32', InceptionTimePlus)\n",
    "InceptionTimePlus47x47 = named_partial('InceptionTimePlus47x47', InceptionTimePlus, nf=47, depth=9)\n",
    "InceptionTimePlus62x62 = named_partial('InceptionTimePlus62x62', InceptionTimePlus, nf=62, depth=9)\n",
    "InceptionTimeXLPlus = named_partial('InceptionTimeXLPlus', InceptionTimePlus, nf=64, depth=12)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tsai.data.core import TSCategorize\n",
    "from tsai.models.utils import count_parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "n_vars = 3\n",
    "seq_len = 51\n",
    "c_out = 2\n",
    "xb = torch.rand(bs, n_vars, seq_len)\n",
    "\n",
    "test_eq(InceptionTimePlus(n_vars,c_out)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out,concat_pool=True)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, bottleneck=False)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, residual=False)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, conv_dropout=.5)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, stoch_depth=.5)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars, c_out, seq_len=seq_len, zero_norm=True, flatten=True)(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, coord=True, separable=True, \n",
    "                          norm='Instance', zero_norm=True, bn_1st=False, fc_dropout=.5, sa=True, se=True, act=nn.PReLU, act_kwargs={})(xb).shape, [bs, c_out])\n",
    "test_eq(InceptionTimePlus(n_vars,c_out, coord=True, separable=True,\n",
    "                          norm='Instance', zero_norm=True, bn_1st=False, act=nn.PReLU, act_kwargs={})(xb).shape, [bs, c_out])\n",
    "test_eq(count_parameters(InceptionTimePlus(3, 2)), 455490)\n",
    "test_eq(count_parameters(InceptionTimePlus(6, 2, **{'coord': True, 'separable': True, 'zero_norm': True})), 77204)\n",
    "test_eq(count_parameters(InceptionTimePlus(3, 2, ks=40)), count_parameters(InceptionTimePlus(3, 2, ks=[9, 19, 39])))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 16\n",
    "n_vars = 3\n",
    "seq_len = 51\n",
    "c_out = 2\n",
    "xb = torch.rand(bs, n_vars, seq_len)\n",
    "\n",
    "model = InceptionTimePlus(n_vars, c_out)\n",
    "model(xb).shape\n",
    "test_eq(model[0](xb), model.backbone(xb))\n",
    "test_eq(model[1](model[0](xb)), model.head(model[0](xb)))\n",
    "test_eq(model[1].state_dict().keys(), model.head.state_dict().keys())\n",
    "test_eq(len(ts_splitter(model)), 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_eq(check_bias(InceptionTimePlus(2,3, zero_norm=True), is_conv)[0].sum(), 0)\n",
    "test_eq(check_weight(InceptionTimePlus(2,3, zero_norm=True), is_bn)[0].sum(), 6)\n",
    "test_eq(check_weight(InceptionTimePlus(2,3), is_bn)[0], np.array([1., 1., 1., 1., 1., 1., 1., 1.]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(10): InceptionTimePlus(n_vars,c_out,stoch_depth=0.8,depth=9,zero_norm=True)(xb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "InceptionTimePlus(\n",
       "  (backbone): Sequential(\n",
       "    (0): InceptionBlockPlus(\n",
       "      (inception): ModuleList(\n",
       "        (0): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(3, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(3, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act): ReLU()\n",
       "        )\n",
       "        (1): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act): ReLU()\n",
       "        )\n",
       "        (2): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (3): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act): ReLU()\n",
       "        )\n",
       "        (4): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "          (act): ReLU()\n",
       "        )\n",
       "        (5): InceptionModulePlus(\n",
       "          (bottleneck): ConvBlock(\n",
       "            (0): AddCoords1d()\n",
       "            (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          )\n",
       "          (convs): ModuleList(\n",
       "            (0): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(39,), stride=(1,), padding=(19,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(19,), stride=(1,), padding=(9,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "            (2): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): SeparableConv1d(\n",
       "                (depthwise_conv): Conv1d(33, 33, kernel_size=(9,), stride=(1,), padding=(4,), groups=33, bias=False)\n",
       "                (pointwise_conv): Conv1d(33, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "              )\n",
       "            )\n",
       "          )\n",
       "          (mp_conv): Sequential(\n",
       "            (0): MaxPool1d(kernel_size=3, stride=1, padding=1, dilation=1, ceil_mode=False)\n",
       "            (1): ConvBlock(\n",
       "              (0): AddCoords1d()\n",
       "              (1): Conv1d(129, 32, kernel_size=(1,), stride=(1,), bias=False)\n",
       "            )\n",
       "          )\n",
       "          (concat): Concat(dim=1)\n",
       "          (norm): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "      )\n",
       "      (shortcut): ModuleList(\n",
       "        (0): ConvBlock(\n",
       "          (0): AddCoords1d()\n",
       "          (1): Conv1d(3, 128, kernel_size=(1,), stride=(1,), bias=False)\n",
       "          (2): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "        )\n",
       "        (1): BatchNorm1d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "      )\n",
       "      (act): ModuleList(\n",
       "        (0): ReLU()\n",
       "        (1): ReLU()\n",
       "      )\n",
       "      (add): Add\n",
       "    )\n",
       "  )\n",
       "  (head): Sequential(\n",
       "    (0): Sequential(\n",
       "      (0): GAP1d(\n",
       "        (gap): AdaptiveAvgPool1d(output_size=1)\n",
       "        (flatten): Reshape(bs)\n",
       "      )\n",
       "      (1): LinBnDrop(\n",
       "        (0): Linear(in_features=128, out_features=3, bias=True)\n",
       "      )\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "net = InceptionTimePlus(2,3,**{'coord': True, 'separable': True, 'zero_norm': True})\n",
    "test_eq(check_weight(net, is_bn)[0], np.array([1., 1., 0., 1., 1., 0., 1., 1.]))\n",
    "net"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|export\n",
    "@delegates(InceptionTimePlus.__init__)\n",
    "class MultiInceptionTimePlus(nn.Sequential):\n",
    "    \"\"\"Class that allows you to create a model with multiple branches of InceptionTimePlus.\"\"\"\n",
    "    _arch = InceptionTimePlus\n",
    "    def __init__(self, feat_list, c_out, seq_len=None, nf=32, nb_filters=None, depth=6, stoch_depth=1., \n",
    "                flatten=False, concat_pool=False, fc_dropout=0., bn=False, y_range=None, custom_head=None, **kwargs):\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            feat_list: list with number of features that will be passed to each body.\n",
    "        \"\"\"\n",
    "        self.feat_list = [feat_list] if isinstance(feat_list, int) else feat_list \n",
    "        \n",
    "        # Backbone\n",
    "        branches = nn.ModuleList()\n",
    "        self.head_nf = 0\n",
    "        for feat in self.feat_list:\n",
    "            if is_listy(feat): feat = len(feat)\n",
    "            m = build_ts_model(self._arch, c_in=feat, c_out=c_out, seq_len=seq_len, nf=nf, nb_filters=nb_filters, \n",
    "                               depth=depth, stoch_depth=stoch_depth, **kwargs)\n",
    "            self.head_nf += output_size_calculator(m[0], feat, ifnone(seq_len, 10))[0]\n",
    "            branches.append(m.backbone)\n",
    "        backbone = _Splitter(self.feat_list, branches)\n",
    "        \n",
    "        # Head\n",
    "        self.c_out = c_out\n",
    "        self.seq_len = seq_len\n",
    "        if custom_head is None:\n",
    "            head = self._arch.create_head(self, self.head_nf, c_out, seq_len, flatten=flatten, concat_pool=concat_pool, \n",
    "                                          fc_dropout=fc_dropout, bn=bn, y_range=y_range)\n",
    "        else: \n",
    "            head = custom_head(self.head_nf, c_out, seq_len)\n",
    "        \n",
    "        layers = OrderedDict([('backbone', nn.Sequential(backbone)), ('head', nn.Sequential(head))])\n",
    "        super().__init__(layers)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#|exporti\n",
    "class _Splitter(Module):\n",
    "    def __init__(self, feat_list, branches):\n",
    "        self.feat_list, self.branches = feat_list, branches\n",
    "    def forward(self, x):\n",
    "        if is_listy(self.feat_list[0]): \n",
    "            x = [x[:, feat] for feat in self.feat_list]\n",
    "        else: \n",
    "            x = torch.split(x, self.feat_list, dim=1)\n",
    "        _out = []\n",
    "        for xi, branch in zip(x, self.branches): _out.append(branch(xi))\n",
    "        output = torch.cat(_out, dim=1)\n",
    "        return output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[W NNPACK.cpp:53] Could not initialize NNPACK! Reason: Unsupported hardware.\n"
     ]
    }
   ],
   "source": [
    "bs = 16\n",
    "n_vars = 3\n",
    "seq_len = 51\n",
    "c_out = 2\n",
    "xb = torch.rand(bs, n_vars, seq_len)\n",
    "\n",
    "test_eq(count_parameters(MultiInceptionTimePlus([1,1,1], c_out)) > count_parameters(MultiInceptionTimePlus(3, c_out)), True)\n",
    "test_eq(MultiInceptionTimePlus([1,1,1], c_out).to(xb.device)(xb).shape, MultiInceptionTimePlus(3, c_out).to(xb.device)(xb).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 5, 2, 10])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): create_conv_lin_nd_head(\n",
       "    (0): Conv1d(128, 10, kernel_size=(1,), stride=(1,))\n",
       "    (1): Linear(in_features=12, out_features=10, bias=True)\n",
       "    (2): Transpose(-1, -2)\n",
       "    (3): Reshape(bs, 5, 2, 10)\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "n_vars = 3\n",
    "seq_len = 12\n",
    "c_out = 10\n",
    "xb = torch.rand(bs, n_vars, seq_len)\n",
    "new_head = partial(conv_lin_nd_head, d=(5,2))\n",
    "net = MultiInceptionTimePlus(n_vars, c_out, seq_len, custom_head=new_head)\n",
    "print(net.to(xb.device)(xb).shape)\n",
    "net.head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([16, 2])\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Sequential(\n",
       "  (0): Sequential(\n",
       "    (0): GAP1d(\n",
       "      (gap): AdaptiveAvgPool1d(output_size=1)\n",
       "      (flatten): Reshape(bs)\n",
       "    )\n",
       "    (1): LinBnDrop(\n",
       "      (0): Linear(in_features=384, out_features=2, bias=True)\n",
       "    )\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": null,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "bs = 16\n",
    "n_vars = 6\n",
    "seq_len = 12\n",
    "c_out = 2\n",
    "xb = torch.rand(bs, n_vars, seq_len)\n",
    "net = MultiInceptionTimePlus([1,2,3], c_out, seq_len)\n",
    "print(net.to(xb.device)(xb).shape)\n",
    "net.head"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "bs = 8\n",
    "c_in = 7  # aka channels, features, variables, dimensions\n",
    "c_out = 2\n",
    "seq_len = 10\n",
    "xb2 = torch.randn(bs, c_in, seq_len)\n",
    "model1 = MultiInceptionTimePlus([2, 5], c_out, seq_len)\n",
    "model2 = MultiInceptionTimePlus([[0,2,5], [0,1,3,4,6]], c_out, seq_len)\n",
    "test_eq(model1.to(xb2.device)(xb2).shape, (bs, c_out))\n",
    "test_eq(model1.to(xb2.device)(xb2).shape, model2.to(xb2.device)(xb2).shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from tsai.data.external import *\n",
    "from tsai.data.core import *\n",
    "from tsai.data.preprocessing import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y, splits = get_UCR_data('NATOPS', split_data=False)\n",
    "tfms  = [None, [TSCategorize()]]\n",
    "batch_tfms = TSStandardize()\n",
    "dls = get_ts_dls(X, y, splits=splits, tfms=tfms, batch_tfms=batch_tfms)\n",
    "model = InceptionTimePlus(dls.vars, dls.c, dls.len)\n",
    "xb,yb=first(dls.train)\n",
    "test_eq(model.to(xb.device)(xb).shape, (dls.bs, dls.c))\n",
    "test_eq(count_parameters(model), 460038)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "X, y, splits = get_UCR_data('NATOPS', split_data=False)\n",
    "tfms  = [None, [TSCategorize()]]\n",
    "batch_tfms = TSStandardize()\n",
    "dls = get_ts_dls(X, y, splits=splits, tfms=tfms, batch_tfms=batch_tfms)\n",
    "model = MultiInceptionTimePlus([4, 15, 5], dls.c, dls.len)\n",
    "xb,yb=first(dls.train)\n",
    "test_eq(model.to(xb.device)(xb).shape, (dls.bs, dls.c))\n",
    "test_eq(count_parameters(model), 1370886)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/javascript": "IPython.notebook.save_checkpoint();",
      "text/plain": [
       "<IPython.core.display.Javascript object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/Users/nacho/notebooks/tsai/nbs/036_models.InceptionTimePlus.ipynb saved at 2023-03-19 14:18:56\n",
      "Correct notebook to script conversion! 😃\n",
      "Sunday 19/03/23 14:18:59 CET\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "                <audio  controls=\"controls\" autoplay=\"autoplay\">\n",
       "                    <source src=\"data:audio/wav;base64,UklGRvQHAABXQVZFZm10IBAAAAABAAEAECcAACBOAAACABAAZGF0YdAHAAAAAPF/iPh/gOoOon6w6ayCoR2ZeyfbjobxK+F2Hs0XjKc5i3DGvzaTlEaraE+zz5uLUl9f46fHpWJdxVSrnfmw8mYEScqUP70cb0Q8X41uysJ1si6Eh1jYzXp9IE2DzOYsftYRyoCY9dJ/8QICgIcEun8D9PmAaBPlfT7lq4MFIlh61tYPiCswIHX+yBaOqT1QbuW7qpVQSv9lu6+xnvRVSlyopAypbGBTUdSalrSTaUBFYpInwUpxOzhti5TOdndyKhCGrdwAfBUcXIJB69p+Vw1egB76+n9q/h6ADglbf4LvnIHfF/981ODThF4m8HiS0riJVjQ6c+/EOZCYQfJrGrhBmPVNMmNArLKhQlkXWYqhbaxXY8ZNHphLuBJsZUEckCTFVHMgNKGJytIDeSUmw4QN4Qx9pReTgb3vYX/TCBuApf75f+P5Y4CRDdN+B+tngk8c8nt03CKGqipgd13OhotwOC5x9MCAknFFcmlmtPmagFFFYOCo0qRzXMhVi57pryNmIEqJlRi8bm52PfuNM8k4dfQv+4cO12l6zCGdg3jl730uE/KAPvS+f0wEAoAsA89/XfXQgBESIn6S5luDtiC8eh/YmIfpLqt1OMp5jXg8/24MveqUNUnPZsqw0Z3yVDldnaUOqIZfXlKrm36zzWhjRhaT+r+ncHI5/otUzfd2uSt7hl/bqXtoHaCC6+mqfrAOeoDD+PJ/xf8RgLMHfH/b8GeBihZIfSXidoQSJWB52NM1iRkzz3MkxpKPbUCrbDu5d5fgTAxkSK3JoEhYD1p2omere2LZTuqYLbdWa49Cx5Dww7tyXDUnioXRkHhwJyKFvd/AfPoYy4Fl7j1/LQorgEr9/X89+0qAOAwAf13sJoL8Gkd8wt25hWIp3Heez/eKODfPcSPCzpFNRDVqf7UlmnNQKGHgqd+jgVvJVm2f265QZTpLS5byur1tpT6ajvrHq3Q2MXWIxtUCehoj8YMk5LB9hRQegeTypn+nBQWA0QHgf7f2q4C5EFt+5ucOg2YfHXtq2SSHpS0ydnTL4IxFO6pvNb4ulBdInWfcsfSc7VMmXpSmE6eeXmZThJxpsgRohEfOk86+AHCoOpOMFsx1dv8s6oYT2k17uR7ngpXod34IEJqAaPfnfyABCIBZBpl/NPI2gTQVjX134x2ExSPMeR7VtYjZMWJ0W8ftjkA/YW1durCWykvjZFKu4p9LVwVbZKNkqpxh6U+6mRC2mGq2Q3SRvsIgcpc2sIpD0Bp4uiiFhW3ecXxOGgaCDe0Vf4cLPoDv+/5/mfw1gN4KKX+17emBqBmYfBHfVYUZKFR44NBtiv41bHJUwx+RJkP1apu2VJlkTwli4qrwoo1ax1dToNCtemRSTBGXz7kJbdM/PY/Dxht0dTLziH7Ul3loJEiE0uJsfdsVTYGL8Yt/AgcMgHYA7X8S+IqAYA+QfjzpxIIVHnp7tdqzhmAstXaxzEqMETpScGC/dJP3Rmdo8LIZnOVSEF+Opxumsl1sVF+dVrE5Z6NIiZSkvVdv2zsqjdnK8HVDLlyHyNjuegogM4NA5z9+YRG9gA722H97AgOA/gSyf43zCIHdE899yuTIg3ciNXpm1jmImTDwdJPITI4RPhRugbvslbFKt2Vfr/6eTFb4W1WkY6m6YPdQjJr2tNZp3EQlko7BgXHRNz2LAc+gdwMq7IUf3R58ohtFgrbr6n7hDFWAlPr8f/T9I4CECU9/De+vgVQY5nxh4POEzybJeCTS5YnCNAZzhsRzkP1Bsmu4t4aYU07nYuerA6KWWcJYO6HHrKJjaE3Zl624UWz/QOOPjcWHc7QzdIk40yl5tCWjhIDhJX0xF4CBMvBsf10IF4Ac//Z/bPlsgAcOwn6S6n6CwxzUewLcRoYaKzV38M23i9o493CNwL6S1UUuaQe0QpvbUfdfiqglpcRccFU+nkWwambASUiVfLyqbg49xY2eyWh1hy/Sh37XjHpaIYKD7OUEfrgS5IC09MV/1gMBgKMDyH/n9N6AhhINfh7mdoMoIZt6r9fAh1cvfHXNya6N4DzDbqi8K5WWSYlmbbAdnkpV6FxJpWSo1V8DUmGb3rMRaQBG2JJgwN9wCDnNi8HNI3dKK1aG0dvHe/UciIJf6rt+Og5wgDn59X9P/xWAKQhxf2XweYH+FjB9suGVhIMlOnlo02GJhTOdc7vFyo/TQGxs2Li7lz9NwmPurBihnVi7WSWiwKvGYntOpJiOt5drKUKMkFnE8HLxNPmJ9NG4eP8mAYUv4Np8hhi3gdruSX+3CSWAwP38f8f6UoCuDPF+6Os8gnAbKnxQ3d2F0imydzDPKIuiN5lxu8EKkrFE82kftW2az1DbYImpMqTUW3FWIJ83r5hl2koJlla7+m0+PmSOZcjcdMgwS4g11iZ6qCLUg5jkxn0QFA6BWvOvfzEFBIBHAtp/Qfa3gC4RSH5y5yeD2B/8evnYS4cULgR2CMsUja47cG/QvW6UeEhXZ3+xP51GVNVdP6Zpp+1eDFM5nMeySWghR4+TNL85cD46YIyCzKJ2kCzEhoTabXtGHs+CCemJfpMPjoDe9+t/qQALgM8Gj3++8UaBqRV2fQTjO4Q3JKd5r9TgiEYyMHTxxiWPpz8jbfq585YpTJpk960xoKFXsVoTo7yq6GGMTw==\" type=\"audio/wav\" />\n",
       "                    Your browser does not support the audio element.\n",
       "                </audio>\n",
       "              "
      ],
      "text/plain": [
       "<IPython.lib.display.Audio object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "#|eval: false\n",
    "#|hide\n",
    "from tsai.export import get_nb_name; nb_name = get_nb_name(locals())\n",
    "from tsai.imports import create_scripts; create_scripts(nb_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "python3",
   "language": "python",
   "name": "python3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
